######################################################################
# estimate model via continuation approach
######################################################################

# activate project environment
cd( dirname(Base.source_path()) )
using Pkg
Pkg.activate("../")

using ArgParse

println("loading NKPH methods")
include("../src/NewKeynesianPreferredHabitatModels.jl");
using .NewKeynesianPreferredHabitatModels


# parse command line arguments
s = ArgParseSettings()
@add_arg_table! s begin
    # input options
    "--model_fname"
        help = "model xlsx filename"
        arg_type = String
        required = true
    "--rates_fname"
        help = "rate data csv filename"
        arg_type = String
        required = true
    "--regs_fname"
        help = "localization regression data csv filename"
        arg_type = String
        required = true
    "--alt_regs_fname"
        help = "alternative localization regression data csv filename"
        arg_type = String
        required = true
    "--rates_date_start"
        help = "start date for rates data"
        arg_type = String
        required = true
    "--rates_date_end"
        help = "end date for rates data"
        arg_type = String
        required = true
    # estimation options
    "--seed"
        help = "seed for random draw of initial parameters"
        arg_type = Int
        default = 1234
    "--draw_init_method"
        help = "method to use for random draw of initial parameters"
        arg_type = Int
        default = 0
    "--draw_init_std"
        help = "variance for random draw of initial parameters (draw_init_method=2)"
        arg_type = Float64
        default = 0.1
    "--N_steps"
        help = "number of steps for continuation estimation"
        arg_type = Int
        default = 11
        # output options
    "--estimates_fname"
        help = "xlsx filename for saving estimates"
        arg_type = String
        default = "output"
end

# get command line arguments
println("parsing options")
parsed_args = parse_args(ARGS, s)

# input options
model_fname = parsed_args["model_fname"]
rates_fname = parsed_args["rates_fname"]
regs_fname = parsed_args["regs_fname"]
alt_regs_fname = parsed_args["alt_regs_fname"]
rates_date_start = parsed_args["rates_date_start"]
rates_date_end = parsed_args["rates_date_end"]
# estimation options
draw_init_method = parsed_args["draw_init_method"]
draw_init_std = parsed_args["draw_init_std"]
seed = parsed_args["seed"]
N_steps = parsed_args["N_steps"]
# output options
# if drawing initial params, append seed to estimates_fname
estimates_fname = parsed_args["estimates_fname"]
if draw_init_method!=0
    estimates_fname *= ("_" * string(seed))
end

println("model:")
println(model_fname)
println("data [rates/regs]:")
println(rates_fname, " ", regs_fname)
println("dates [rates]:")
println(rates_date_start, " ", rates_date_end)
println("intial param options:")
println(draw_init_method, " ", draw_init_std, " ", seed)
println("continuation steps:")
println(N_steps)
println("saving estimates to:")
println(estimates_fname)
flush(stdout)


# setup model
println("initializing model...")
model_fpath = MODEL_DIR * model_fname * ".xlsx"
rates_fpath = MOMENT_DIR * rates_fname * ".csv"
regs_fpath = MOMENT_DIR * regs_fname * ".csv"
alt_regs_fpath = MOMENT_DIR * alt_regs_fname * ".csv"

# create estimation object
nkphopt = NKPHEstimationOptimizer(
    model_fpath, rates_fpath, regs_fpath, alt_regs_fpath, 
    rates_date_start, rates_date_end
)
# initialize model objects
initialize_empirical_moments!(nkphopt)
initialize_model!(nkphopt)
initialize_targets!(nkphopt)
setup_nlopt!(nkphopt)

# initial values
println("setting up initial targets")
p0 = draw_init_params(seed, nkphopt, draw_init_method, draw_init_std)
set_continuation_targets!(p0, nkphopt)

# display initial model info
println("initial value results:")
println("="^100); println("="^100);
println("max root: ", maximum(abs.(nkphopt.nkphsolvr.root_vec)) )
println("min/max Upsilon eigenvals: ", nkphopt.nkphtgts.Upsilon_min_pos_eigval[1],
    " ", nkphopt.nkphtgts.Upsilon_max_neg_eigval[1])
println("min M eigenval: ", nkphopt.nkphtgts.M_min_eigval[1])
loss_val = nkphopt.nkphtgts.loss_val[1]
println("loss value: ", loss_val)


println("="^100); println("="^100); flush(stdout)
println("beginning continuation estimation")
println("="^100); println("="^100); flush(stdout)
loss_val, x_opt, ret = estimate_model_continuation!(p0, N_steps, nkphopt)

# save
println("saving")
save_estimates(loss_val, x_opt, estimates_fname, nkphopt)
println("="^100); println("="^100)
